import os
import json
import numpy as np
from typing import List, Dict, Tuple, Any
from rainforeLearn.gomoku.v2.configs.config import GomokuDQNConfig


class ExpertDataLoader:
    """专家棋谱加载器 - 负责加载和预处理专家棋谱数据"""

    def __init__(self, config: GomokuDQNConfig):
        self.config = config
        self.board_size = config.board_size
        self.expert_games = []
        self.current_index = 0

        if config.use_expert_data:
            self._load_expert_data()

    def _load_expert_data(self) -> None:
        """加载专家棋谱数据"""
        expert_path = self.config.expert_data_path
        if not os.path.exists(expert_path):
            print(f"⚠️ 专家棋谱路径不存在: {expert_path}")
            return

        # 遍历文件夹中的所有JSON文件
        for filename in os.listdir(expert_path):
            if filename.endswith('.json'):
                file_path = os.path.join(expert_path, filename)
                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        data = json.load(f)
                        self._process_expert_file(data)
                except Exception as e:
                    print(f"❌ 加载专家棋谱文件失败: {file_path}, 错误: {e}")

        print(f"✅ 成功加载 {len(self.expert_games)} 局专家棋谱")

    def _process_expert_file(self, data: Dict) -> None:
        """处理单个专家棋谱文件"""
        # 检查棋盘大小是否匹配
        if data.get('metadata', {}).get('board_size') != self.board_size:
            print(f"⚠️ 棋谱棋盘大小 {data.get('metadata', {}).get('board_size')} 与配置不匹配 {self.board_size}")
            return

        # 添加游戏到列表
        for game in data.get('games', []):
            processed_game = self._convert_game_to_transitions(game)
            if processed_game:
                self.expert_games.append(processed_game)

                # 如果启用数据增强，添加旋转和翻转版本
                if self.config.expert_data_augmentation:
                    augmented_games = self._augment_game(processed_game)
                    self.expert_games.extend(augmented_games)

    def _convert_game_to_transitions(self, game: Dict) -> List[Tuple]:
        """将游戏转换为transitions列表"""
        transitions = []
        board = np.zeros((self.board_size, self.board_size), dtype=np.int8)
        moves = game.get('moves', [])
        result = game.get('result', 0)

        for i, move in enumerate(moves):
            player = move.get('player')
            row, col = move.get('position')
            action = row * self.board_size + col

            # 保存当前状态
            state = board.flatten().copy()

            # 执行动作
            board[row, col] = player

            # 计算奖励
            reward = self._calculate_expert_reward(i, len(moves), player, result)

            # 检查是否游戏结束
            done = (i == len(moves) - 1)

            # 保存下一个状态
            next_state = board.flatten().copy()

            # 创建transition
            transition = (state, action, reward, next_state, done)
            transitions.append(transition)

        return transitions

    def _calculate_expert_reward(self, move_idx: int, total_moves: int,
                                 player: int, result: int) -> float:
        """计算专家棋谱中的奖励值"""
        # 基础奖励
        if result == player:  # 获胜方
            reward = 0.1
        elif result == 0:  # 平局
            reward = 0.0
        else:  # 失败方
            reward = -0.05

        # 最后一步特殊奖励
        if move_idx == total_moves - 1:
            if result == player:  # 获胜方的最后一步
                reward = 1.0
            elif result == -player:  # 失败方的最后一步
                reward = -0.5

        # 应用奖励缩放
        return reward * self.config.expert_reward_scale

    def _augment_game(self, game: List[Tuple]) -> List[List[Tuple]]:
        """对游戏数据进行增强（旋转和翻转）"""
        augmented_games = []

        # 实现旋转和翻转逻辑
        # 这里需要对每个transition中的state, action和next_state进行变换
        # 可以实现90度、180度、270度旋转以及水平、垂直翻转

        return augmented_games

    def sample_batch(self, batch_size: int) -> List[Tuple]:
        """采样一批专家数据"""
        if not self.expert_games:
            return []

        # 根据采样策略选择游戏
        if self.config.expert_sampling_strategy == 'uniform':
            # 均匀采样
            game_indices = np.random.choice(len(self.expert_games), batch_size, replace=True)
            batch = []
            for idx in game_indices:
                game = self.expert_games[idx]
                # 从游戏中随机选择一个transition
                if game:
                    transition_idx = np.random.randint(0, len(game))
                    batch.append(game[transition_idx])
            return batch
        else:
            # 其他采样策略实现
            pass
